//
// Ephi - simulation of magnetic fields and particles
// Copyright (C) 2007 Indrek Mandre <indrek(at)mare.ee>
// For more information please see http://www.mare.ee/indrek/ephi/
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
//

#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#include "ephi.hpp"

#define WIDTH 800
#define HEIGHT 800

#define BEGIN_SPEED 15.0
#define END_SPEED 100000.0

#define COILMATRIX_SIDE 8

enum eshape { ONE, ROUND, SQUARE };
enum eaim { CORNER, OPEN };

static void usage (const char *appname)
{
  fprintf (stderr, "%s: %s [-a corner|open] [-c one|round|square] [-s spacing] [-i currentN:startCurrent:stepSize] [-r radiusN:startRadius:stepSize] [-w wireRadiusN:startWireRadius:stepSize] [-g]\n", appname, appname);
  exit(1);
}

static bool make_picture = false;

struct mytask : Task
{
  prec_t current;
  prec_t radius;
  prec_t wr;
  prec_t spacing;
  eaim aim;
  prec_t end_speed;
  ShapedCoilFactory *shaper;

  void execute ()
  {
    Screen *screen1 = 0;
    Screen *screen2 = 0;
    end_speed = 0;
    prec_t rr = radius + (wr * 2 + spacing) / PREC_SQRT2;
    vect3d pos(0, 0, 0);
    vect3d dir;
    switch (aim)
      {
        case CORNER:
          dir = vect3d(-rr, radius, 0);
          break;
        case OPEN:
          dir = vect3d(-rr, radius/PREC_SQRT2, radius/PREC_SQRT2);
          break;
      }

    if ( make_picture )
      {
        screen1 = new Screen (WIDTH, HEIGHT, wr * 6, wr * 6);
        screen1->set_axis (dir, vect3d(1, 0, 0), vect3d(0, 1, 0));
        screen1->draw_axes ();
        screen2 = new Screen (WIDTH, HEIGHT, wr * 6, wr * 6);
        screen2->set_axis (dir, vect3d(1, 0, 0), vect3d(0, 0, 1));
        screen2->draw_axes ();
      }

    Statics statics;
    make_polywell_cube (statics, radius, wr, spacing, 0, 0); // empty shell
    make_polywell_cube (statics, radius, wr, spacing, current, 0, *shaper);

    if ( make_picture )
      {
        statics.draw (*screen1);
        statics.draw (*screen2);
      }

    prec_t speed_begin = ElectroDynamics::eV2ms (BEGIN_SPEED, ELECTRON_MASS);
    prec_t speed_end = ElectroDynamics::eV2ms (END_SPEED, ELECTRON_MASS);
    prec_t speed = 0;
    if ( speed == 0 || speed >= END_SPEED )
        speed = (speed_begin + speed_end) / 2;

    while (1)
      {
        ElectroDynamics dyn(statics);
        dyn.setDTM (0.1, 5);
        //dyn.setDTM (0.0001, 1); // speed test
        size_t n = dyn.inject_electron (pos, dir, ElectroDynamics::ms2eV (speed, ELECTRON_MASS));
        size_t ticks = dyn.ticks_for_distance (speed, radius * 2);

        while (--ticks)
          {
            dyn.step ();
            if ( make_picture )
              {
                screen1->set_pixel (dyn.get_pos(n), Screen::RED);
                screen2->set_pixel (dyn.get_pos(n), Screen::RED);
              }
            if ( dyn.inContact (dyn.get_pos(n), n) )
              {
                if ( make_picture )
                  {
                    screen1->draw_cross (dyn.get_pos(n), Screen::GREEN, 0.1 * wr);
                    screen2->draw_cross (dyn.get_pos(n), Screen::GREEN, 0.1 * wr);
                  }
                break;
              }
          }

        if ( !ticks )
            speed_begin = speed;
        else
            speed_end = speed;

        if ( speed_end - speed_begin < 500 ) // with a 0.5 km/s precision
            break;

        speed = (speed_begin + speed_end)/2;
      }

    end_speed = (speed_begin + speed_end)/2;
    if ( speed_end == END_SPEED )
      {
        printf ("current=%f, radius=%f, wr=%f, spacing=%f, pierce speed=X\n",
            prec2double(current), prec2double(radius), prec2double(wr),
            prec2double(spacing));
      }
    else
      {
        printf ("current=%f, radius=%f, wr=%f, spacing=%f, pierce speed=%.0f eV (%.1f km/s)\n",
            prec2double(current), prec2double(radius), prec2double(wr),
            prec2double(spacing), prec2double(ElectroDynamics::ms2eV(end_speed, ELECTRON_MASS)),
            prec2double(end_speed/1000.0));
      }

    if ( make_picture )
      {
        screen1->write ("polywell_pierce1.bmp");
        delete screen1;
        screen2->write ("polywell_pierce2.bmp");
        delete screen2;
      }
  }
};

int main (int argc, char *argv[])
{
  eshape shape = ONE;
  eaim aim = CORNER;
  bool scale_i_with_wr = false;
  int ICOUNT = 16;
  double ISTART = 100.0;
  double IGROW = 100.0;
  int RCOUNT = 6;
  double RSTART = 0.1;
  double RGROW = 0.05;
  int WCOUNT = 1;
  double WSTART = 0.035;
  double WGROW = 0.005;
  double SPACING = 0.01;

  for ( int i = 1; i < argc; i++ )
    {
      if ( argv[i][0] != '-' )
          usage(argv[0]);

      switch (argv[i][1])
        {
          case 'a': {
            if ( ++i == argc )
                usage(argv[0]);
            if ( strcmp (argv[i], "corner") == 0 )
                aim = CORNER;
            else if ( strcmp (argv[i], "open") == 0 )
                aim = OPEN;
            else
                usage(argv[0]);
            break;
          }
          case 'c': {
            if ( ++i == argc )
                usage(argv[0]);
            if ( strcmp (argv[i], "round") == 0 )
                shape = ROUND;
            else if ( strcmp (argv[i], "square") == 0 )
                shape = SQUARE;
            else if ( strcmp (argv[i], "one") == 0 )
                shape = ONE;
            else
                usage(argv[0]);
            break;
          }
          case 's': {
            if ( ++i == argc )
                usage(argv[0]);
            SPACING = atof (argv[i]);
            break;
          }
          case 'i': {
            if ( ++i == argc )
                usage(argv[0]);
            if ( sscanf (argv[i], "%d:%lf:%lf", &ICOUNT, &ISTART, &IGROW) != 3 )
                usage(argv[0]);
            break;
          }
          case 'r': {
            if ( ++i == argc )
                usage(argv[0]);
            if ( sscanf (argv[i], "%d:%lf:%lf", &RCOUNT, &RSTART, &RGROW) != 3 )
                usage(argv[0]);
            break;
          }
          case 'w': {
            if ( ++i == argc )
                usage(argv[0]);
            if ( sscanf (argv[i], "%d:%lf:%lf", &WCOUNT, &WSTART, &WGROW) != 3 )
                usage(argv[0]);
            break;
          }
          case 'g': {
            scale_i_with_wr = true;
            break;
          }
          case 'p': {
            make_picture = true;
            break;
          }
          default:
            usage(argv[0]);
            break;
        }
    }

  ShapedCoilFactory *shaper = 0;

  switch (shape)
    {
      case ONE:
        shaper = new ShapedCoilFactory();
        break;
      case ROUND:
        shaper = new ShapedCoilFactory (COILMATRIX_SIDE, COILMATRIX_SIDE, 1/PREC_SQRT2, 1/PREC_SQRT2);
        shaper->set_circle();
        break;
      case SQUARE:
        shaper = new ShapedCoilFactory (COILMATRIX_SIDE, COILMATRIX_SIDE, 1/PREC_SQRT2, 1/PREC_SQRT2);
        break;
    }

  Ephi::setDebugLevel (Ephi::DEBUG_ERROR);

  TaskManager tmgr;
  std::vector<mytask*> tasks;
  for ( int radiusN = 0; radiusN < RCOUNT; radiusN++ )
    {
      for ( int wN = 0; wN < WCOUNT; wN++ )
        {
          for ( int currentN = 0; currentN < ICOUNT; currentN++ )
            {
              if ( make_picture && !tasks.empty() ) // only one needed when making a picture
                  break;
              mytask *task = new mytask ();
              task->shaper = shaper;
              task->radius = RSTART + radiusN * RGROW;
              task->wr = WSTART + wN * WGROW;
              if ( scale_i_with_wr )
                  task->current = ISTART * (task->wr * task->wr) / (WSTART * WSTART);
              else
                  task->current = ISTART + currentN * IGROW;
              task->spacing = SPACING;
              task->aim = aim;
              tasks.push_back (task);
              tmgr.addTask (task);
            }
        }
    }

  while (tmgr.run());
  delete shaper;

  FILE *fp = fopen ("polywell_pierce.csv", "w+");
  if ( !fp )
    {
      Ephi::debug (Ephi::DEBUG_CRITICAL, "unable to open polywell_pierce.csv for writing");
      exit(-1);
    }

  fprintf (fp, "RADIUS[m],WR[m],SPACING[m],CURRENT[A],ENERGY[eV],SPEED[km/s]\n");
  for ( size_t i = 0; i < tasks.size(); i++ )
    {
      fprintf (fp, "%f,%f,%f,%f,%.0f,%.1f\n",
          prec2double(tasks[i]->radius), prec2double(tasks[i]->wr), prec2double(tasks[i]->spacing),
          prec2double(tasks[i]->current), prec2double(ElectroDynamics::ms2eV(tasks[i]->end_speed, ELECTRON_MASS)),
          prec2double(tasks[i]->end_speed/1000.0));
      delete tasks[i];
    }
  fclose (fp);

  return 0;
}

